"""
backend_api.py
只负责提供 /api/* 路由，使用 Flask Blueprint
"""

import os
import logging
import torch
import numpy as np
from flask import Blueprint, request, jsonify
from flask_cors import CORS
from neo4j import GraphDatabase, exceptions
from transformers import BertModel, BertTokenizer

# ------------------ 日志 ------------------
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# ------------------ 创建 Blueprint ------------------
bp = Blueprint('api', __name__)
CORS(bp, resources={r"/api/*": {"origins": "*"}})

# ------------------ 配置 ------------------
NEO4J_URI      = "bolt://localhost:7687"
NEO4J_USER     = "neo4j"
NEO4J_PASSWORD = "gjj030712"

def create_neo4j_driver():
    try:
        driver = GraphDatabase.driver(
            NEO4J_URI,
            auth=(NEO4J_USER, NEO4J_PASSWORD),
            max_connection_lifetime=3600
        )
        with driver.session() as session:
            session.run("RETURN 1")
        logging.info("✅ Neo4j 连接成功")
        return driver
    except Exception as e:
        logging.error(f"❌ Neo4j 连接失败: {e}")
        return None

driver = create_neo4j_driver()

# ------------------ BERT ------------------
# ✅ 仅修改这一行：指向新的英文路径
MODEL_DIR = r"D:\bert_models"

tokenizer = BertTokenizer.from_pretrained(MODEL_DIR)
bert = BertModel.from_pretrained(MODEL_DIR)
device = "cuda" if torch.cuda.is_available() else "cpu"
bert.to(device).eval()
logging.info(f"✅ BERT 加载完成，设备：{device}")

@torch.no_grad()
def encode_sentences(sentences):
    encoded = tokenizer(
        sentences,
        padding=True,
        truncation=True,
        return_tensors="pt"
    ).to(device)
    return torch.nn.functional.normalize(
        bert(**encoded).last_hidden_state[:, 0], dim=-1
    )

# ------------------ 缓存人物向量 ------------------
ALL_CHARS, CHAR_VECS = [], None
if driver:
    with driver.session() as session:
        ALL_CHARS = [r["name"] for r in session.run("MATCH (p:Person) RETURN p.name AS name")]
    if ALL_CHARS:
        CHAR_VECS = encode_sentences(ALL_CHARS).cpu().numpy()
        logging.info(f"📦 已缓存 {len(ALL_CHARS)} 个人物的向量")

# ------------------ 工具函数 ------------------
def get_character_relationships(name, depth=1):
    if not driver:
        return {"error": "数据库连接失败"}, 500
    try:
        with driver.session() as session:
            query = f"""
            MATCH (source:Person)-[rels*1..{depth}]-(target)
            WHERE source.name = $name
            UNWIND rels AS rel
            RETURN DISTINCT
                   source.name AS source_name,
                   target.name AS target_name,
                   rel.type   AS relationship_type
            """
            result = session.run(query, name=name)
            nodes, links = set(), []
            nodes.add(name)
            for rec in result:
                src, tgt, rel = rec["source_name"], rec["target_name"], rec["relationship_type"]
                nodes.update([src, tgt])
                links.append({
                    "id": f"{min(src, tgt)}-{max(src, tgt)}-{rel}",
                    "source": src,
                    "target": tgt,
                    "type": rel
                })
            nodes = [
                {"id": n, "name": n,
                 "connections": sum(1 for l in links if l["source"] == n or l["target"] == n)}
                for n in nodes
            ]
            return {"nodes": nodes, "links": links}, 200
    except exceptions.Neo4jError as e:
        logging.error(f"查询关系时出错: {e}")
        return {"error": str(e)}, 500

# ------------------ 路由 ------------------
@bp.route('/health', methods=['GET'])
def health_check():
    global driver
    if not driver:
        driver = create_neo4j_driver()
    if not driver:
        return jsonify({"status": "unhealthy", "database": "not initialized"}), 500
    try:
        with driver.session() as session:
            session.run("RETURN 1")
        return jsonify({"status": "healthy", "database": "connected"}), 200
    except Exception:
        driver = create_neo4j_driver()
        return jsonify({"status": "unhealthy", "database": "disconnected"}), 500

@bp.route('/characters', methods=['GET'])
def get_characters():
    global driver
    if not driver:
        driver = create_neo4j_driver()
    if not driver:
        return jsonify({"error": "数据库连接失败"}), 500
    try:
        limit = max(1, min(int(request.args.get('limit', 20)), 100))
        with driver.session() as session:
            res = session.run(
                "MATCH (p:Person) RETURN p.name AS name ORDER BY name LIMIT $limit",
                {"limit": limit}
            )
            return jsonify([{"name": r["name"]} for r in res])
    except Exception as e:
        logging.error(f"获取人物列表时出错: {e}")
        return jsonify({"error": "服务器内部错误"}), 500

@bp.route('/characters/search', methods=['GET'])
def search_characters():
    if not driver:
        return jsonify([])
    q = request.args.get('q', '')
    if not q:
        return jsonify([])
    try:
        with driver.session() as session:
            res = session.run(
                "MATCH (p:Person) WHERE p.name CONTAINS $query RETURN p.name AS name ORDER BY p.name LIMIT 10",
                {"query": q}
            )
            return jsonify([{"name": r["name"]} for r in res])
    except exceptions.Neo4jError as e:
        logging.error(f"搜索人物时出错: {e}")
        return jsonify([])

@bp.route('/characters/<name>/relationships', methods=['GET'])
def get_relationships(name):
    depth = min(max(request.args.get('depth', 1, int), 1), 4)
    data, status = get_character_relationships(name, depth)
    return jsonify(data), status

@bp.route('/characters/<name>', methods=['GET'])
def get_character_details(name):
    if not driver:
        return jsonify({"error": "数据库连接失败"}), 500
    try:
        with driver.session() as session:
            rec = session.run("""
                MATCH (p:Person) WHERE p.name = $name
                OPTIONAL MATCH (p)-[r]->()
                RETURN p.name AS name,
                       COUNT(r) AS relationship_count,
                       COLLECT(DISTINCT type(r)) AS relationship_types
            """, name=name).single()
            if not rec:
                return jsonify({"error": "人物不存在"}), 404
            return jsonify({
                "name": rec["name"],
                "relationship_count": rec["relationship_count"],
                "relationship_types": rec["relationship_types"]
            })
    except exceptions.Neo4jError as e:
        logging.error(f"获取人物详情时出错: {e}")
        return jsonify({"error": str(e)}), 500

@bp.route("/search_semantic", methods=["GET"])
def search_semantic():
    if CHAR_VECS is None:
        return jsonify([])
    q = request.args.get("q", "").strip()
    if not q:
        return jsonify([])
    try:
        q_vec = encode_sentences([q]).cpu().numpy()
        sims = (CHAR_VECS @ q_vec.T).squeeze()
        topk = min(10, len(ALL_CHARS))
        idx = np.argsort(sims)[-topk:][::-1]
        return jsonify([
            {"name": ALL_CHARS[i], "score": float(sims[i])}
            for i in idx
        ])
    except Exception as e:
        logging.error(f"语义搜索出错: {e}")
        return jsonify([])